# -*- coding: UTF-8 -*-
import sys, os, re, glob
from pprint import pprint
import copy
from concurrent.futures import ThreadPoolExecutor,as_completed
from pysubparser import parser
from subprocess import Popen, PIPE

#FIXME 图片宽度，暂时不可自定义
imgwidth = 350

#FIXME PDF页面CSS，暂时不可自定义
#tr { page-break-inside: avoid !important; }
style = '''
<style>
/* DivTable.com */
.divTable{
	display: table;
    height: 29.0cm;
    table-layout: fixed;
}
.divTableRow {
	display: table-row;
}
.imgRow {
    height: 20%;
}
.divTableCell {
	display: table-cell;
    width: 50%;
}
.divTableCellDash {
  padding-left: 5px;
  border-left: 1px dashed #999999;
  border-collapse: collapse;
}
.divTableBody {
	display: table-row-group;
}

.text {
    vertical-align: top;
    line-height: 110%;
    text-align: left;
    font-size: 28px;
    letter-spacing: 1px;
    font-family: "Times New Roman", Times, serif;
}

/* Page Breaks */
.pb_before {
    page-break-before: always !important;
 	font-size: 50px;
}

.pb_after {
    page-break-after: always !important;
}

/***Avoid page break before the element (if possible)***/
.pb_before_avoid {
    page-break-before: avoid !important;
}

/***Avoid page break after the element (if possible)***/
.pb_after_avoid {
    page-break-after: avoid !important;
}

/* Avoid page break inside the element (if possible) */
.pbi_avoid {
    page-break-inside: avoid !important;
}

</style>
'''
#文件名合格处理
def ChangeFileName(filename):
    if filename == None:
        return None
    filename = filename.replace('\\','')\
    .replace('/','')\
    .replace('*','')\
    .replace('<','')\
    .replace('>','')\
    .replace('|','')\
    .replace('?','？')\
    .replace(':','：')\
    .replace('"','')\
    .replace('\t','')\
    .replace('\n','')
    while filename.endswith('.'):
        filename = filename[:-1]
    filename = filename.strip()
    return filename

def stamptosec(ts):
    tss = [float(tx) for tx in ts.split(':')]
    sec = tss[0] * 60 * 60
    sec+= tss[1] * 60
    sec+= tss[2]
    return sec

class FFBook():
    def __init__(self):
        self.obj = None
        self.srcRoot = ''
        self.outRoot = os.path.join(os.path.dirname(os.path.abspath(__file__))\
                , 'out')
        self.tasks   = 1
        self.sub_type = ['ass', 'ssa', 'srt', 'sub', 'txt', 'vtt']
        self.media_type = ['mkv', 'mp4', 'avi']
        self.media_sub  = []
        pass

    def createHtml(self, htmlf, mobif, subs):
        imgpath = os.path.dirname(htmlf)
        #title = subs[0]['text'][0]
        htitle = os.path.basename(imgpath)
        #if len(htitle) > 75:
        #    htitle = htitle[:70] + '...'
        #'''
        htmlstr = f'<html xmlns="http://www.w3.org/1999/xhtml"><head>\
        <meta http-equiv="content-type" content="text/html; charset=utf-8" />\
        {style}<title>{htitle}</title></head><body>'
        sbnum = 0
        sblen = len(subs)
        emptyimg = '<img src="" width="0" height="0" alt="" />'
        #for mobi
        mobistr = f'<html xmlns="http://www.w3.org/1999/xhtml"><head>\
        <meta http-equiv="content-type" content="text/html; charset=utf-8" />\
        {style}<title>{htitle}</title></head><body>'
        for si, sub in enumerate(subs):
            if si > 0:
                imgf = os.path.join(imgpath, '{:04d}.png'.format(si))
                mobistr += f'<img src="file:///{imgf}" /><br />'

            subtxt = ' '.join(sub['text'])
            subtxt = re.sub(r'<font color="#[^"]+">([^<]*)</font>', '\\1', subtxt)
            mobistr += subtxt + '<br />'

        mobistr += '</body></html>'
        with open(mobif, "w", encoding='utf-8') as f:
            #print(mobistr)
            f.write(mobistr)

        #for mobi
        while sbnum < sblen:
            if sbnum == 0:
                imgf1 = emptyimg
            else:
                img1 = os.path.join(imgpath, '{:04d}.png'.format(sbnum))
                imgf1 = f'<img width="{imgwidth}" src="file:///{img1}">'

            subtxt1 = ' '.join(subs[sbnum]['text'])
            subtxt1 = re.sub(r'<font color="#[^"]+">([^<]*)</font>', '\\1', subtxt1)
            pgnum1  = sbnum + 1

            if sbnum + 1 < sblen:
                img2 = os.path.join(imgpath, '{:04d}.png'.format(sbnum + 1))
                imgf2 = f'<img width="{imgwidth}" src="file:///{img2}">'
                subtxt2 = ' '.join(subs[sbnum + 1]['text'])
                subtxt2 = re.sub(r'<font color="#[^"]+">([^<]*)</font>', '\\1', subtxt2)
                pgnum2  = sbnum + 2
            else:
                imgf2 = emptyimg
                subtxt2 = ''
                pgnum2  = ''
            if sbnum + 2 < sblen:
                img3 = os.path.join(imgpath, '{:04d}.png'.format(sbnum + 2))
                imgf3 = f'<img width="{imgwidth}" src="file:///{img3}">'
                subtxt3 = ' '.join(subs[sbnum + 2]['text'])
                subtxt3 = re.sub(r'<font color="#[^"]+">([^<]*)</font>', '\\1', subtxt3)
                pgnum3 = sbnum + 3
            else:
                imgf3 = emptyimg
                subtxt3 = ''
                pgnum3 = ''
            if sbnum + 3 < sblen:
                img4 = os.path.join(imgpath, '{:04d}.png'.format(sbnum + 3))
                imgf4 = f'<img width="{imgwidth}" src="file:///{img4}">'
                subtxt4 = ' '.join(subs[sbnum + 3]['text'])
                subtxt4 = re.sub(r'<font color="#[^"]+">([^<]*)</font>', '\\1', subtxt4)
                pgnum4 = sbnum + 4
            else:
                imgf4 = emptyimg
                subtxt4 = ''
                pgnum4 = ''


            htmlstr += f'''
            <div class="divTable pb_before pbi_avoid">
                <div class="divTableBody">
					<div class="divTableRow">
						<div class="divTableCell imgRow">{imgf1}</div>
						<div class="divTableCell divTableCellDash imgRow">{imgf2}</div>
					</div>
					<div class="divTableRow">
						<div class="divTableCell text">{subtxt1}</div>
						<div class="divTableCell text divTableCellDash">{subtxt2}</div>
					</div>
					<div class="divTableRow">
						<div class="divTableCell imgRow">{imgf3}</div>
						<div class="divTableCell divTableCellDash imgRow">{imgf4}</div>
					</div>
					<div class="divTableRow">
						<div class="divTableCell text">{subtxt3}</div>
						<div class="divTableCell text divTableCellDash">{subtxt4}</div>
					</div>
                </div>
            </div>
            '''

            #cls = ''
            #if sbnum + 4 > sblen:
            #    cls = 'class="last-page"'
            #htmlstr += f'''
            #<page size="A4" {cls}>
            #<table style="width:98%;height:100%">
            #    <tr>
            #      <th width="49%" height="20%">{imgf1}</th>
            #      <!---<th rowspan=2 class="rpgnum">{pgnum1}</th>
            #      <th rowspan=2 class="lpgnum bleft">{pgnum2}</th>--->
            #      <th width="49%" class="bleft">{imgf2}</th>
            #    </tr>
            #    <tr>
            #      <td class="text" height="30%">{subtxt1}</td>
            #      <td class="text bleft" height="30%">{subtxt2}</td>
            #    </tr>
            #    <tr>
            #      <th width="49%" height="20%" class="btop">{imgf3}</th>
            #      <!---<th rowspan=2 class="rpgnum btop">{pgnum3}</th>
            #      <th rowspan=2 class="lpgnum bleft btop">{pgnum4}</th>--->
            #      <th width="49%" class="bleft">{imgf4}</th>
            #    </tr>
            #    <tr>
            #      <td class="text" height="30%">{subtxt3}</td>
            #      <td class="text bleft" height="30%">{subtxt4}</td>
            #    </tr>
            #  </table>
            #  </page>
            #  '''

            sbnum += 4

        htmlstr += '</body></html>'
        with open(htmlf, "w", encoding='utf-8') as f:
             f.write(htmlstr)
        #'''

    def createPdf(self, pdf, htmlist):
        #toc = {
        #    'xsl-style-sheet': 'toc.xsl'
        #}
        try:
            cmd = 'wkhtmltopdf \
                --page-size A4 \
                --minimum-font-size 20 \
                --enable-local-file-access \
                --disable-smart-shrinking \
                --margin-top 0mm \
                --margin-bottom 5mm \
                --margin-left 0mm \
                --margin-right 0mm \
                --encoding UTF-8 \
                --outline-depth 10 \
                --footer-font-size 13'

            if sys.platform == 'win32':
                cmd += ' --footer-center [page]'
            else:
                cmd += " --footer-center '[page]'"
            htmlfs = ''
            for htm in sorted(htmlist):
                htmlfs += '"file:///' + htm + '" '
            #print(htmlfs)
            cmd += f' toc --xsl-style-sheet toc.xsl {htmlfs} "{pdf}"'
            #print(cmd)
            ret = self.run_command(cmd)
            if ret == 0:
                if self.obj:
                    self.obj.gmsg(f'PDF: {pdf}制作成功！')
            else:
                if self.obj:
                    self.obj.cmsg(f'PDF: {pdf}制作失败！:{cmd}')
            #pdfkit.from_file(htmlist, pdf, options = options, toc = toc)
        except:
            if self.obj:
                self.obj.cmsg(f'PDF: {pdf}制作失败！:{cmd}')
        return ret



    def createMobi(self, mobi, htmloist, dpname):
        #mobi
        index_html = f'<html xmlns="http://www.w3.org/1999/xhtml"><head><meta http-equiv="content-type" content="text/html; charset=utf-8" /><title>Table of Contents</title></head><body><h2 align="center">Table of Contents</h2><ul>'
        for htmlf in sorted(htmloist):
            bshname = os.path.basename(htmlf).replace('.mobi.html', '')
            index_html += f'<li><a href="file:///{htmlf}">{bshname}</li>'

        index_html += '</ul></body></html>'
        index_file = os.path.join(os.path.dirname(htmlf), 'index.html')
        with open(index_file, 'w', encoding='utf-8') as f:
             f.write(index_html)


        #cover_img = os.path.join(tmp_dir, level, dpname, 'cover.png')
        cmd_line = f'ebook-convert "{index_file}" "{mobi}" --language=en \
        --authors ginger --use-auto-toc --toc-threshold=0 --toc-title="Contents" \
        --mobi-file-type both --title="{dpname}"'
        #print(cmd_line)
        #p = Popen(cmd_line, shell=True, stdout=PIPE)
        #p.wait()
        ret = self.run_command(cmd_line)
        if ret == 0 and self.obj:
            self.obj.gmsg(f'电子书{mobi}制作成功！')
        else:
            if self.obj:
                self.obj.cmsg(f'电子书{mobi}制作失败：{cmd_line}！')
        return ret

    def keep_sentence_short(self, word):
        # 标题按空格拆分，发现超长单词增加空格
        ftt = ''
        for ft in word.split():
            if len(ft) > 35:
                ft = ft[:35] + ' ' + ft[35:]
            ftt += ' ' + ft

        # 防止生成过长的文件名导致toc标题过长
        if len(ftt) > 75:
            ftt = ftt[:70] + '...'
        return ChangeFileName(ftt)


    def convert(self, outD, medias, mtitle):
        media = medias['media']
        subtitlef = medias['subtitle']
        subtitles = parser.parse(subtitlef, encoding='utf-8')
        jt = 0
        sec = ''
        sec2 = ''
        asec = ''
        mtext = f'<font size="50" face="Times-Roman" color=black>§.</font> <h5 \
        style="display: inline;">{mtitle.strip()}</h5>'
        subl = [{'text':[mtext], 'start': 0, 'end': 0}]
        tsmin = 0
        temax = 0
        ss = {}
        subtext = []
        pg = 4
        tile = '2x2'
        imgheight = int(250 * imgwidth / 400)
        scale= f'{imgwidth}:{imgheight}'
        for sn, subtitle in enumerate(subtitles):
            ts = stamptosec(str(subtitle.start))
            if sn == 0:
                tsmin = ts
            te = stamptosec(str(subtitle.end))
            tt = ts + (te - ts) / 2
            sec += 'lt(prev_pts*TB\,{})*gte(pts*TB\,{})+'.format(round(tt, 1), round(tt, 1))
            if sn != 0 and sn % pg == 0:
                ss['text'] = copy.copy(subtext)
                ss['start'] = tsmin
                ss['end'] = temax
                subl.append(copy.deepcopy(ss))
                tt2 = ss['start'] + (ss['end'] - ss['start']) / 2
                sec2 += 'lt(prev_pts*TB\,{})*gte(pts*TB\,{})+'.format(round(tt2, 1), round(tt2, 1))
                asec += f'{tsmin},'

                subtext.clear()
                ss.clear()
                tsmin = ts
                subtext.append(subtitle.text)
            else:
                #print(sn, sn % pg)
                subtext.append(subtitle.text)
                tsmin = ts if tsmin > ts else tsmin
                temax = te if te > temax else temax
            jt = sn

        #最后4句别漏
        ss['text'] = copy.copy(subtext)
        ss['start'] = tsmin
        ss['end'] = temax
        subl.append(copy.deepcopy(ss))
        tt2 = ss['start'] + (ss['end'] - ss['start']) / 2
        sec2 += 'lt(prev_pts*TB\,{})*gte(pts*TB\,{})+'.format(round(tt2, 1), round(tt2, 1))
        asec += f'{tsmin},'

        sec = sec[:-1]
        sec2 = sec2[:-1]
        asec = asec[:-1]
        if sys.platform == 'win32':
            #FIXME Wired!!!
            subtitlef = subtitlef.replace('/', '\\')\
                    .replace('\\', '\\\\\\\\')\
                    .replace(':', '\\\\\\:')\
                    .replace("'", "\\\\\\'")\
                    .replace(",", "\\\\\\,")\
                    .replace("[", "\\\\\\[")\
                    .replace("]", "\\\\\\]")
        else:
            #FIXME Wired!!!
            subtitlef = subtitlef.replace("'", "\\\\\\\\\\\'")\
                                 .replace(",", "\\\\\\\\\\\,")\
                                 .replace("[", "\\\\\\\\\\[")\
                                 .replace("]", "\\\\\\\\\\]")


        #sty = f'-vf "select=\'{sec2}\',subtitles={subtitlef},scale={scale}" -vsync vfr'
        #ff = FFmpeg(
        #    #inputs={media: '-hide_banner -loglevel panic -accurate_seek -copyts'},
        #    inputs={media: '-hide_banner -accurate_seek -copyts'},
        #    outputs={f"{outD}{os.sep}%04d.png": sty}
        #)
        #split audio for oid
        #acmd = 'ffmpeg -hide_banner -accurate_seek -i "{mp4}" -vn -f segment \
        #        -segment_times {asec} "{out}{sep}%04d.mp3"'\
        #.format(mp4=media, out=outD, sep=os.sep, asec=asec)
              #"select=\'{sec2}\',scale={scale},subtitles={subtitlef},tile={tile}" \
        cmd = f'ffmpeg -accurate_seek -copyts -i "{media}" \
              -vf \
              "select=\'{sec2}\',scale={scale},subtitles={subtitlef}" \
              -vsync vfr "{outD}{os.sep}%04d.png"'

        img0 = os.path.join(outD, '0001.png')
        #print(img0, cmd)
        if not os.path.exists(img0):
            #print(img0)
            ret = self.run_command(cmd)
            if ret == 0 and self.obj:
                self.obj.gmsg(f'{media} 成功截图')
            else:
                if self.obj:
                    self.obj.cmsg(f'{media} 截图失败：{cmd}')

        htmlf = os.path.join(outD, f'{mtitle}.html')
        mobif = os.path.join(outD, f'{mtitle}.mobi.html')
        #if not os.path.exists(htmlf) or not os.path.exists(mobif):
        self.createHtml(htmlf, mobif, subl)


        return '{} 处理完成，共{}个截图'.format(media, jt + 1)


    def run_command(self, command):
        process = Popen(command, shell=True, \
                stdout=PIPE, stderr=PIPE, \
                stdin=PIPE)
        #print(command)
        output, err = process.communicate()
        #print(output.strip())
        process.stdin.close()
        #while True:
        #    #output = process.stdout.readline()
        #    #process.stderr.close()
        #    #print(type(output), f'[{output}]')
        #    if output == b'' and process.poll() is not None:
        #        break
        #    if output:
        #        print(output.strip())
        #        #self.obj.gmsg(output.strip())

        rc = process.poll()
        return rc


    def check_sub(self):
        g = os.walk(self.srcRoot)
        for path,dir_list,file_list in g:
            for file_name in file_list:
                media = os.path.join(path, file_name)
                #print(path, '         ', file_name)
                media_match = re.match(r'(.*)\.(mkv|avi|mp4)$', media,
                        re.IGNORECASE)
                if media_match:
                    media_name = media_match.group(1)
                    sflag = 0
                    for subt in self.sub_type:
                        subf = os.path.join(glob.escape(path),
                                glob.escape(f'{media_name}') + '*' + f'.{subt}')
                        gsubf = glob.glob(subf)
                        if len(gsubf) > 0:
                            subtitlef = os.path.join(path, gsubf[0])
                            if not subtitlef.endswith('.vtt'):
                                self.media_sub.append({'media': media, 'subtitle':
                                    subtitlef })
                                sflag = 1
                            else:
                                sflag = 2
                            break


                    if sflag != 1:
                        # 生成的字幕文件名
                        msub = os.path.join(path, f'{media_name}.srt')
                        if sflag == 2:
                            print(f'{file_name}存在Vtt字幕，需要转换......')
                            cmd = f'ffmpeg -hide_banner -loglevel panic -y -i \
                            "{subtitlef}" "{msub}"'
                        else:
                            print(f'{file_name}不存在现成字幕，尝试从视频中获取......')
                            cmd = f'ffmpeg -hide_banner -loglevel panic -y -i \
                            "{media}" -map 0:s:0 "{msub}"'
                        ret = self.run_command(cmd)
                        if ret == 0:
                            if self.obj:
                                self.obj.gmsg(f'{file_name}字幕获取成功')
                            self.media_sub.append({'media': media, 'subtitle':
                                msub })
                        else:
                            if self.obj:
                                self.obj.cmsg(f'{file_name}字幕获取失败:{cmd}')

    def main_task(self, obj):
        if len(self.media_sub) == 0:
            print('没有发现可以处理的文件!')
            return None

        self.obj = obj
        tt = []
        htmlist = []
        htmloist = []
        tpool = ThreadPoolExecutor(max_workers=int(self.tasks))
        for medias in self.media_sub:
            media_name = os.path.splitext(os.path.basename(medias['media']))[0]
            # 防止生成过长的文件名导致toc标题过长
            mtitle = self.keep_sentence_short(media_name)
            outD = os.path.join(self.outRoot,
                os.path.basename(self.srcRoot), media_name)
            if not os.path.exists(outD):
                os.makedirs(outD)
            f_t = tpool.submit(self.convert, outD, medias, mtitle)
            tt.append(f_t)
            htmlist.append(os.path.join(outD , f'{mtitle}.html'))
            htmloist.append(os.path.join(outD, f'{mtitle}.mobi.html'))

        for fu in as_completed(tt):
            result = fu.result()
            if result != None:
                pass


        # 开始生成pdf, mobi
        output_name = ChangeFileName(os.path.basename(self.srcRoot))
        pdf = os.path.join(self.outRoot, output_name + '.pdf')
        pret = self.createPdf(pdf, htmlist)
        mobi = os.path.join(self.outRoot, output_name + '.mobi')
        mret = self.createMobi(mobi, htmloist, output_name)
        if pret == 0 and mret == 0 and os.path.exists(pdf) and os.path.exists(mobi):
            return [pdf, mobi]

    def set_conf(self, srcRoot, outRoot, tasks):
        #print(f'处理目录: {srcRoot}, 任务数: {tasks} , 输出路径: {outRoot}')
        self.srcRoot = srcRoot
        self.outRoot = outRoot
        self.tasks = tasks


if __name__ == '__main__':
    ffbook = FFBook()
    # 读取配置文件
    tasks, outRoot, srcRoot = 1, os.path.abspath('./out'), sys.argv[1]
    ffbook.set_conf(srcRoot, outRoot, tasks)
    print(f'处理目录: {srcRoot}, 任务数: {tasks} , 输出路径: {outRoot}')

    # 检查目录是否有字幕文件，没有则尝试从视频文件获取
    ffbook.check_sub()
    # 开始执行主要任务
    ffbook.main_task(None)

